/*
 * @Author: January
 * @Date: 2021-12-22 21:16:47
 */
#include <libecc/lib_ecc_config.h>
#include <libecc/hash/hash_algs.h>
#include <libecc/sig/sig_algs.h>
#ifdef CONFIG_WITH_LOG
//#define CONFIG_DEBUG
#include <stdio.h>
#include <internal/log.h>
#else
#define LOG_DEBUG(fmt, ...)
#define LOG_DEBUG_VERBOSE(fmt, ...)
#define LOG_DEBUG_MEM(fmt, ...)
#endif

typedef struct {
    prj_pt pt;
    uint8_t* data;
    int data_len;
    uint8_t* hash;
    int hash_len;
} sm2_ciphertext_t;

/**
 * sm2_ciphertext ::= 
 *      x       INTEGER
 *      y       INTEGER
 *      data    OCTSTRING
 *      hash    OCTSTRING
 */

#ifdef WITH_ENC_SM2

//TODO: 不属于sm2的功能，后续应该移动到单独的文件中
#include <libecc/encode/base64.h>
int pem2der(uint8_t* pem, int inlen, uint8_t* der, int out_buffer_len) {
    int ret = -1;
    uint8_t pem_buffer[256];
    int p_pem_buffer = 0;
    int der_len = 0;

    if (inlen > 256) {
        LOG_DEBUG("data too long");
        goto out;
    }
    for(int i=0; i<inlen; i++) {
        if(pem[i] == '-') {
            while(pem[i] != '\n' && i<inlen) {
                i++;
            }
        } else {
            while (pem[i] != '\n' && i<inlen) {
                pem_buffer[p_pem_buffer] = pem[i];
                p_pem_buffer++;
                i++;
            }
        }
    }
    ret = base64_decode(der, out_buffer_len, &der_len, pem_buffer, p_pem_buffer);
    if (ret < 0) {
        LOG_DEBUG("Fail to decode pem");
        goto out;
    }
    ret = der_len;
    out:
    return ret;
}

void x963_kdf(const hash_mapping* hash, const void *in, size_t inlen, void* out, size_t outlen) {
    uint32_t counter = 1;
    uint32_t counter_be;
    uint8_t dgst[SM3_DIGEST_SIZE];
    uint32_t dgstlen = SM3_DIGEST_SIZE;
    uint8_t *pout = out;
    size_t rlen = outlen;
    size_t len;
    hash_context ctx;

    while (rlen > 0) {
#if __BYTE_ORDER__ ==__ORDER_LITTLE_ENDIAN__
        counter_be = __builtin_bswap32(counter);
#else
        counter_be = counter;
#endif
        counter++;
        hash->hfunc_init(&ctx);
        hash->hfunc_update(&ctx, in, inlen);
        hash->hfunc_update(&ctx, (uint8_t*)&counter_be, sizeof(counter_be));
        hash->hfunc_finalize(&ctx, dgst);
        len = (dgstlen <= rlen) ? dgstlen : rlen;
        local_memcpy(pout, dgst, len);
        rlen -= len;
        pout += len;
    }
}

static int sm2_do_encrypt(const hash_mapping* hash, const unsigned char *in,
                          int inlen,  ec_pub_key* pubkey, sm2_ciphertext_t *ciphertext) {
    int ret = -1;
    const ec_params* ec = pubkey->params;

    int hash_size = hash->digest_size;
    prj_pt share_point;
    prj_pt_t ephem_point = &(ciphertext->pt);
    nn k;
    unsigned char buf[64];
    size_t i;
    hash_context ctx;

    /* check arguments */
    if (!in || !pubkey) {
        LOG_DEBUG_VERBOSE("invalid input argument ");
        goto out;
    }

    if ((uint32_t)inlen > ciphertext->data_len
    || (uint32_t)hash_size > ciphertext->hash_len) {
        LOG_DEBUG_VERBOSE("invalid plain text length");
        goto out;
    }
    prj_pt_mul(&share_point, &(pubkey->params->ec_gen_cofactor), &(pubkey->y));
    if (prj_pt_iszero(&share_point)) {
        LOG_DEBUG_VERBOSE("invalid public key");
        goto out;
    }

    /* rand k in [1, n-1] */
    do {
        ret = nn_get_random_mod(&k, &(ec->ec_gen_order));
        if (ret < 0) {
            LOG_DEBUG_VERBOSE("Fail to generate random");
            goto out;
        }
    } while(nn_iszero(&k));

    /* compute ephem_point [k]G = (x1, y1) */
    prj_pt_mul(ephem_point, &k, &(ec->ec_gen));
    // ephem_point is the pointer to ciphertext->c1 and since
    // c1 in the ciphertext has been generated, do not modify it any longer.

    /* compute ECDH share_point [k]P_B = (x2, y2) */
    prj_pt_mul(&share_point, &k, &(pubkey->y));

    /* compute t = KDF(x2 || y2, klen) */
    ret = prj_pt_export_to_aff_buf(&share_point, buf, 64);
    if (ret < 0) {
        goto out;
    }
    x963_kdf(hash, buf, 64, ciphertext->data, inlen);

    /* ciphertext = t xor in */
    for (i = 0; i < inlen; i++) {
        ciphertext->data[i] ^= in[i];
    }
    ciphertext->data_len = inlen;
    /* generate hash = Hash(x2 || M || y2) */

    hash->hfunc_init(&ctx);
    hash->hfunc_update(&ctx, buf, hash_size);
    hash->hfunc_update(&ctx, in, inlen);
    hash->hfunc_update(&ctx, buf + hash_size, hash_size);
    hash->hfunc_finalize(&ctx, ciphertext->hash);
    ciphertext->hash_len = hash_size;
    ret = 0;
    out:
    return ret;
}

static int sm2_do_decrypt(const hash_mapping* hash, ec_priv_key* privkey, const sm2_ciphertext_t *ciphertext,
                          unsigned char *out, int out_buf_len) {
    int ret = -1;
    const ec_params* ec = privkey->params;

    int hash_size = hash->digest_size;
    prj_pt share_point;
    prj_pt_src_t ephem_point = &(ciphertext->pt);
    unsigned char buf[64];
    size_t i;
    hash_context ctx;

    /* check arguments */
    if (!out || !privkey) {
        LOG_DEBUG_VERBOSE("invalid input argument ");
        goto out;
    }

    if ((uint32_t)ciphertext->data_len > out_buf_len
        || (uint32_t)ciphertext->hash_len != hash_size ) {
        LOG_DEBUG_VERBOSE("invalid plain text length");
        goto out;
    }
    prj_pt_mul(&share_point, &(ec->ec_gen_cofactor), ephem_point);
    if (prj_pt_iszero(&share_point)) {
        LOG_DEBUG_VERBOSE("invalid ciphertext");
        goto out;
    }

    /* compute ECDH share_point [d]C1 = (x2, y2) */
    prj_pt_mul(&share_point, &(privkey->x), ephem_point);

    /* compute t = KDF(x2 || y2, klen) */
    ret = prj_pt_export_to_aff_buf(&share_point, buf, 64);
    if (ret < 0) {
        goto out;
    }
    x963_kdf(hash, buf, 64, out, ciphertext->data_len);

    /* ciphertext = t xor in */
    for (i = 0; i < ciphertext->data_len; i++) {
        out[i] ^= ciphertext->data[i];
    }
    /* generate hash = Hash(x2 || M || y2) */
    hash->hfunc_init(&ctx);
    hash->hfunc_update(&ctx, buf, hash_size);
    hash->hfunc_update(&ctx, out, ciphertext->data_len);
    hash->hfunc_update(&ctx, buf + hash_size, hash_size);
    hash->hfunc_finalize(&ctx, buf);
    if (local_memcmp(buf, ciphertext->hash, ciphertext->hash_len) != 0) {
        ret = -1;
        goto out;
    }
    ret = ciphertext->data_len;
    out:
    return ret;
}

int sm2_import_subpubkey(ec_pub_key* pubkey, ec_params* ec, uint8_t* in, int inlen) {
    int ret = -1;

    if (inlen != 91) {
        LOG_DEBUG("Invalid sm2 subpubkey");
        goto out;
    }
    ret = ec_pub_key_import_from_aff_buf(pubkey, ec, in + 27, 64, SM2);
    out:
    return ret;
}

int sm2_import_sec1privkey(ec_key_pair* keypair, ec_params* ec, uint8_t* in, int inlen) {
    int ret = -1;

    if (inlen < 39 || inlen > 121) {
        LOG_DEBUG("Invalid sm2 sec1privkey");
        goto out;
    }

    ec_priv_key_import_from_buf(&(keypair->priv_key), ec, in + 7, 32, SM2);
    LOG_DEBUG_MEM(in+7, 32, "priv key:");
    if (inlen == 121) {
        LOG_DEBUG_MEM(in + 57, 64, "pub key:");
        ret = ec_pub_key_import_from_aff_buf(&(keypair->pub_key), ec, in + 57, 64, SM2);
    } else {
        ret = 0;
    }
    out:
    return ret;
}

int sm2_convert_openssl_sig(uint8_t* in, int inlen, uint8_t* out, int out_buffer_len) {
    int ret = -1;
    uint8_t* p_len_r;
    uint8_t* p_len_s;

    if (inlen < 70 || out_buffer_len < 64) {
        LOG_DEBUG_VERBOSE("Invalid sig len");
        goto out;
    }

    p_len_r = in + 3;
    if (*p_len_r == 32) {
        local_memcpy(out, p_len_r + 1, 32);
    } else if(*p_len_r == 33) {
        local_memcpy(out, p_len_r + 2, 32);
    } else {
        LOG_DEBUG_VERBOSE("Invalid [r] size: %d", *p_len_r);
        ret = -1;
        goto out;
    }

    p_len_s = p_len_r + 1 + *p_len_r + 1;
    if (*p_len_s == 32) {
        local_memcpy(out + 32, p_len_s + 1, 32);
    } else if(*p_len_s == 33) {
        local_memcpy(out + 32, p_len_s + 2, 32);
    } else {
        LOG_DEBUG_VERBOSE("Invalid [s] size: %d", *p_len_s);
        ret = -1;
        goto out;
    }
    ret = 64;
    out:
    return ret;
}

int sm2_convert_openssl_enc(uint8_t* in, int inlen, uint8_t* out, int out_buffer_len) {
    int ret = -1;
    int total_len = 0;
    if (inlen < 70) {
        goto out;
    }

    uint8_t* p_c1_x;
    uint8_t* p_c1_y;
    uint8_t* p_data;
    uint8_t* p_hash_len;
    uint8_t* p_data_len;

    if (inlen < 70 || out_buffer_len < 64) {
        LOG_DEBUG_VERBOSE("Invalid input data");
        goto out;
    }

    p_c1_x = in + 1;
    if (*p_c1_x & 0x80) {
        p_c1_x += (*p_c1_x) ^ 0x80;
        p_c1_x += 1;
    } else {
        p_c1_x += 1;
    }
    p_c1_x++;
    if (*p_c1_x == 32) {
        local_memcpy(out, p_c1_x + 1, 32);
    } else if(*p_c1_x == 33) {
        local_memcpy(out, p_c1_x + 2, 32);
    } else {
        LOG_DEBUG_VERBOSE("Invalid input data");
        ret = -1;
        goto out;
    }
    total_len += 32;

    p_c1_y = p_c1_x + 1 + *p_c1_x + 1;
    if (*p_c1_y == 32) {
        local_memcpy(out + total_len, p_c1_y + 1, 32);
    } else if(*p_c1_y == 33) {
        local_memcpy(out + total_len, p_c1_y + 2, 32);
    } else {
        LOG_DEBUG_VERBOSE("Invalid input data");
        ret = -1;
        goto out;
    }
    total_len += 32;

    p_hash_len = p_c1_y + 1 + *p_c1_y + 1;
    if ((in + inlen < p_hash_len + *p_hash_len)
        || (total_len + *p_hash_len > out_buffer_len) ) {
        LOG_DEBUG_VERBOSE("Invalid input data");
        goto out;
    }

    // TODO: 仅支持小于256字节的数据
    p_data_len = p_hash_len + 1 + *p_hash_len + 1;
    if (*p_data_len & 0x80) {
        p_data = (*p_data_len ^ 0x80) + p_data_len + 1;
        p_data_len += 1;
    } else {
        p_data = p_data_len + 1;
    }
    if ((in + inlen < p_data_len + *p_data_len)
        || (total_len + *p_data_len + *p_hash_len > out_buffer_len) ) {
        LOG_DEBUG_VERBOSE("Invalid input data");
        goto out;
    }
    local_memcpy(out + total_len, p_data, *p_data_len);
    total_len += *p_data_len;
    local_memcpy(out + total_len, p_hash_len + 1, *p_hash_len);
    total_len += *p_hash_len;
    ret = total_len;
    out:
    return ret;
}

int sm2_encrypt(hash_alg_type hash_alg, ec_pub_key* pubkey, const uint8_t* in, int inlen, uint8_t* out, int out_buf_len) {
    int ret = -1;
    sm2_ciphertext_t ciphertext;
    uint8_t hash_size = 0;
    const hash_mapping* hash= NULL;

    hash = get_hash_by_type(hash_alg);
    if (hash == NULL) {
        goto out;
    }

    if (get_hash_sizes(hash_alg, &hash_size, NULL) < 0) {
        goto out;
    }

    if (out_buf_len < (inlen + 64 + hash_size)) {
        goto out;
    }
    // 使用out中的空间为ciphertext分配内存
    ciphertext.data = out + 64;
    ciphertext.data_len = inlen;
    ciphertext.hash = ciphertext.data + inlen;
    ciphertext.hash_len = hash_size;

    ret = sm2_do_encrypt(hash, in, inlen, pubkey, &ciphertext);
    if (ret < 0) {
       goto out;
    }
    prj_pt_export_to_aff_buf(&(ciphertext.pt), out, 64);
    ret = inlen + 64 + hash_size;
    out:
    return ret;
}

int sm2_decrypt(hash_alg_type hash_alg, ec_priv_key* privkey, const uint8_t* in, int inlen, uint8_t* out, int out_buf_len) {
    int ret = -1;
    uint8_t hash_size = 0;
    sm2_ciphertext_t ciphertext;
    const hash_mapping* hash= NULL;

    hash = get_hash_by_type(hash_alg);
    if (hash == NULL) {
        goto out;
    }

    if (get_hash_sizes(hash_alg, &hash_size, NULL) < 0) {
        goto out;
    }
    if (inlen <= 64 + hash_size) {
        LOG_DEBUG_VERBOSE("Invalid input");
        goto out;
    }

    if (out_buf_len < inlen - 64 - hash_size) {
        LOG_DEBUG_VERBOSE("Out buffer too small");
        ret = -1;
    }
    // 使用in中的数据填充ciphertext结构
    ciphertext.data = in + 64;
    ciphertext.data_len = inlen - 64 - hash_size;
    ciphertext.hash = ciphertext.data + ciphertext.data_len;
    ciphertext.hash_len = hash_size;
    prj_pt_import_from_aff_buf(&(ciphertext.pt), in, 64, &(privkey->params->ec_curve));
    ret = sm2_do_decrypt(hash, privkey, &ciphertext, out, out_buf_len);
    out:
    return ret;
}

#endif